package com.example.mybatisplus.threadPool;

import com.example.mybatisplus.holder.ClientAttribute;
import com.example.mybatisplus.holder.ClientAttributeHolder;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.util.concurrent.ListenableFuture;

import java.util.concurrent.Callable;
import java.util.concurrent.Future;


/**
 * 重写线程池，传递线程变量
 */
public class ThreadPoolExecutors extends ThreadPoolTaskExecutor {

	@Override
	public void execute(final Runnable task) {
		final String clientAttribute = ClientAttributeHolder.get();
		super.execute(() -> execute(task, clientAttribute));
	}

	@Override
	public void execute(final Runnable task, long startTimeout) {
		final String clientAttribute = ClientAttributeHolder.get();
		super.execute(() -> execute(task, clientAttribute), startTimeout);
	}

	@Override
	public Future<?> submit(final Runnable task) {
		final String clientAttribute = ClientAttributeHolder.get();
		return super.submit(() -> execute(task, clientAttribute));
	}

	@Override
	public <T> Future<T> submit(final Callable<T> task) {
		final String clientAttribute = ClientAttributeHolder.get();
		return super.submit(() -> execute(task, clientAttribute));
	}

	@Override
	public ListenableFuture<?> submitListenable(final Runnable task) {
		final String clientAttribute = ClientAttributeHolder.get();
		return super.submitListenable(() -> execute(task, clientAttribute));
	}

	@Override
	public <T> ListenableFuture<T> submitListenable(final Callable<T> task) {
		final String clientAttribute = ClientAttributeHolder.get();
		return super.submitListenable(() -> execute(task,clientAttribute));
	}

	private void execute(Runnable task, String clientAttribute) {
		ClientAttributeHolder.set(clientAttribute);
		task.run();
		ClientAttributeHolder.remove();
	}

	private <T> T execute(Callable<T> task, String clientAttribute) throws Exception {
		ClientAttributeHolder.set(clientAttribute);
		T value = task.call();
		ClientAttributeHolder.remove();
		return value;
	}




}
